// pkg/ssh.go - SSH连接模块
package pkg

import (
	"fmt"
	"sync"
	"time"
	"log"
	"golang.org/x/crypto/ssh"
)

// ExecuteCommand 执行SSH命令
func (rct *RemoteControlTool) ExecuteCommand(host HostInfo, command string, useRoot bool) (string, error) {
	config := &ssh.ClientConfig{
		User: host.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(host.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         30 * time.Second,
	}

	addr := fmt.Sprintf("%s:%d", host.IP, host.Port)
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return "", fmt.Errorf("连接失败: %v", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("创建会话失败: %v", err)
	}
	defer session.Close()

	var finalCommand string
	if useRoot && host.RootPass != "" {
		finalCommand = fmt.Sprintf(`echo '%s' | su - -c '%s'`, host.RootPass, command)
	} else {
		finalCommand = command
	}

	output, err := session.CombinedOutput(finalCommand)
	if err != nil {
		return string(output), fmt.Errorf("命令执行失败: %v", err)
	}

	return string(output), nil
}

// ExecuteOnSelectedHosts 在选中的主机上执行命令
func (rct *RemoteControlTool) ExecuteOnSelectedHosts(hostIDs []int, command string, useRoot bool) []TaskResult {
	hosts := make([]HostInfo, 0, len(hostIDs))
	for _, id := range hostIDs {
		host, err := rct.GetHostByID(id)
		if err != nil {
			log.Printf("获取主机ID %d 失败: %v", id, err)
			continue
		}
		hosts = append(hosts, *host)
	}

	if len(hosts) == 0 {
		fmt.Println("没有找到有效的主机")
		return nil
	}

	var results []TaskResult
	var wg sync.WaitGroup
	var mutex sync.Mutex

	for _, host := range hosts {
		wg.Add(1)
		go func(h HostInfo) {
			defer wg.Done()
			output, err := rct.ExecuteCommand(h, command, useRoot)
			result := TaskResult{
				IP:        h.IP,
				Command:   command,
				Output:    output,
				Timestamp: time.Now().Format("2006-01-02 15:04:05"),
			}
			if err != nil {
				result.Error = err.Error()
				result.Success = false
			} else {
				result.Success = true
			}
			mutex.Lock()
			results = append(results, result)
			rct.logToFile(result)
			mutex.Unlock()
		}(host)
	}

	wg.Wait()
	return results
}

// ExecuteMultipleCommandsOnSelectedHosts 在选中的主机上执行多个命令
func (rct *RemoteControlTool) ExecuteMultipleCommandsOnSelectedHosts(hostIDs []int, commands []string, useRoot bool) [][]TaskResult {
	var allResults [][]TaskResult
	for _, command := range commands {
		results := rct.ExecuteOnSelectedHosts(hostIDs, command, useRoot)
		allResults = append(allResults, results)
		rct.logSummaryToFile(results, command, useRoot)
	}
	return allResults
}


